#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高级股票筛选器 - 第二步执行模块
根据投资条件筛选候选股票
"""

import sys
import os
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from real_data_fetcher import RealDataFetcher

class AdvancedStockScreener:
    """高级股票筛选器"""
    
    def __init__(self):
        self.data_fetcher = RealDataFetcher()
        
        # 投资筛选条件
        self.screening_criteria = {
            'min_market_cap': 50,      # 最小市值50亿
            'max_market_cap': 2000,    # 最大市值2000亿
            'min_price': 4,            # 最低价格4元
            'max_price': 20,           # 最高价格20元
            'min_pe_ratio': 1,         # 最小市盈率1 (盈利)
            'max_pe_ratio': 50,        # 最大市盈率50
            'prefer_dividend': True,   # 优先有分红
            'min_volume': 1000000,     # 最小成交量100万
        }
        
        # 精选的已知优质股票候选池
        self.quality_candidates = [
            {'code': '000001', 'name': '平安银行'},
            {'code': '000002', 'name': '万科A'},
            {'code': '600036', 'name': '招商银行'},
            {'code': '002415', 'name': '海康威视'},
            {'code': '002142', 'name': '宁波银行'},
            {'code': '600309', 'name': '万华化学'},
            {'code': '000858', 'name': '五粮液'},
            {'code': '600900', 'name': '长江电力'},
            {'code': '002594', 'name': '比亚迪'},
            {'code': '000725', 'name': '京东方A'},
            {'code': '002916', 'name': '深南电路'},
            {'code': '601318', 'name': '中国平安'},
            {'code': '000063', 'name': '中兴通讯'},
            {'code': '002230', 'name': '科大讯飞'},
            {'code': '002304', 'name': '洋河股份'},
        ]
        
    def print_header(self, title):
        """打印标题"""
        print(f"\n{'='*80}")
        print(f"📊 {title}")
        print(f"{'='*80}")
        
    def get_stock_basic_info(self, stock_code):
        """获取股票基本信息"""
        try:
            # 获取基本信息
            basic_info = self.data_fetcher.get_stock_basic_info(stock_code)
            if basic_info:
                return {
                    'code': stock_code,
                    'name': basic_info.get('name', f'股票{stock_code}'),
                    'current_price': basic_info.get('close', 0),  # 使用close字段
                    'market_cap': basic_info.get('market_cap', 0),
                    'pe_ratio': basic_info.get('pe_ratio', 0),
                    'volume': basic_info.get('volume', 1000000),  # 默认成交量
                    'turnover_rate': basic_info.get('turnover_rate', 0),
                }
        except Exception as e:
            print(f"获取 {stock_code} 基本信息失败: {e}")
            return None
            
    def check_dividend_history(self, stock_code):
        """检查分红历史"""
        try:
            dividend_data = self.data_fetcher.get_dividend_info(stock_code)
            if dividend_data is not None and not dividend_data.empty:
                # 近3年有分红记录
                recent_dividends = dividend_data[dividend_data['announce_date'] >= '2021-01-01'] if 'announce_date' in dividend_data.columns else dividend_data
                return len(recent_dividends) > 0, len(dividend_data)
            return False, 0
        except Exception as e:
            print(f"获取 {stock_code} 分红信息失败: {e}")
            return False, 0
            
    def apply_screening_criteria(self, stock_info):
        """应用筛选条件"""
        if not stock_info:
            return False, "基本信息缺失"
            
        criteria = self.screening_criteria
        
        # 检查市值
        market_cap = stock_info.get('market_cap', 0)
        if market_cap < criteria['min_market_cap'] or market_cap > criteria['max_market_cap']:
            return False, f"市值不符合({market_cap:.1f}亿)"
            
        # 检查股价
        price = stock_info.get('current_price', 0)
        if price < criteria['min_price'] or price > criteria['max_price']:
            return False, f"股价不符合({price:.2f}元)"
            
        # 检查市盈率
        pe_ratio = stock_info.get('pe_ratio', 0)
        if pe_ratio <= 0 or pe_ratio < criteria['min_pe_ratio'] or pe_ratio > criteria['max_pe_ratio']:
            return False, f"市盈率不符合({pe_ratio:.1f})"
            
        # 检查成交量
        volume = stock_info.get('volume', 0)
        if volume < criteria['min_volume']:
            return False, f"成交量不足({volume})"
            
        return True, "符合条件"
        
    def calculate_screening_score(self, stock_info, has_dividend, dividend_count):
        """计算筛选评分"""
        score = 0
        
        # 基础分数
        score += 20
        
        # 市值评分 (优先中等市值)
        market_cap = stock_info.get('market_cap', 0)
        if 100 <= market_cap <= 500:
            score += 25
        elif 50 <= market_cap <= 100 or 500 <= market_cap <= 1000:
            score += 20
        else:
            score += 10
            
        # 价格评分 (优先中等价格)
        price = stock_info.get('current_price', 0)
        if 8 <= price <= 15:
            score += 20
        elif 6 <= price <= 8 or 15 <= price <= 18:
            score += 15
        else:
            score += 10
            
        # 市盈率评分 (优先较低PE)
        pe_ratio = stock_info.get('pe_ratio', 0)
        if 5 <= pe_ratio <= 15:
            score += 20
        elif 15 <= pe_ratio <= 25:
            score += 15
        else:
            score += 10
            
        # 分红评分
        if has_dividend:
            score += 15
            if dividend_count >= 5:
                score += 10
            elif dividend_count >= 3:
                score += 5
                
        return min(score, 100)  # 最高100分
        
    def screen_candidate_stocks(self):
        """筛选候选股票"""
        self.print_header("第二步：股票筛选系统 - 候选股票筛选")
        
        print(f"📋 筛选条件:")
        print(f"   • 市值: {self.screening_criteria['min_market_cap']}-{self.screening_criteria['max_market_cap']}亿")
        print(f"   • 股价: {self.screening_criteria['min_price']}-{self.screening_criteria['max_price']}元")
        print(f"   • 市盈率: {self.screening_criteria['min_pe_ratio']}-{self.screening_criteria['max_pe_ratio']}")
        print(f"   • 优先有分红股票")
        print(f"   • 最小成交量: {self.screening_criteria['min_volume']:,}")
        
        screened_stocks = []
        
        print(f"\n🔍 开始筛选 {len(self.quality_candidates)} 只候选股票...")
        
        for i, candidate in enumerate(self.quality_candidates, 1):
            stock_code = candidate['code']
            stock_name = candidate['name']
            
            print(f"\n[{i:2d}/{len(self.quality_candidates)}] 分析 {stock_code} {stock_name}")
            
            # 获取基本信息
            stock_info = self.get_stock_basic_info(stock_code)
            if not stock_info:
                print(f"   ❌ 无法获取基本信息")
                continue
                
            # 应用筛选条件
            passed, reason = self.apply_screening_criteria(stock_info)
            
            # 检查分红历史
            has_dividend, dividend_count = self.check_dividend_history(stock_code)
            
            if passed:
                # 计算评分
                score = self.calculate_screening_score(stock_info, has_dividend, dividend_count)
                
                stock_result = {
                    **stock_info,
                    'has_dividend': has_dividend,
                    'dividend_count': dividend_count,
                    'screening_score': score
                }
                
                screened_stocks.append(stock_result)
                
                print(f"   ✅ 通过筛选 - 评分: {score:.1f}")
                print(f"      价格: {stock_info['current_price']:.2f}元, 市值: {stock_info['market_cap']:.1f}亿")
                print(f"      PE: {stock_info['pe_ratio']:.1f}, 分红: {'是' if has_dividend else '否'}({dividend_count}次)")
                
            else:
                print(f"   ❌ 未通过筛选 - {reason}")
                
        # 按评分排序
        screened_stocks.sort(key=lambda x: x['screening_score'], reverse=True)
        
        # 生成筛选结果报告
        self.generate_screening_report(screened_stocks)
        
        return screened_stocks
        
    def generate_screening_report(self, screened_stocks):
        """生成筛选结果报告"""
        print(f"\n📊 筛选结果统计:")
        print(f"   • 候选股票数量: {len(self.quality_candidates)}")
        print(f"   • 通过筛选数量: {len(screened_stocks)}")
        print(f"   • 筛选通过率: {len(screened_stocks)/len(self.quality_candidates)*100:.1f}%")
        
        if len(screened_stocks) > 0:
            print(f"\n🏆 筛选结果TOP10:")
            print(f"{'排名':<4} {'代码':<8} {'名称':<12} {'价格':<8} {'市值':<8} {'PE':<6} {'分红':<6} {'评分':<6}")
            print("-" * 70)
            
            for i, stock in enumerate(screened_stocks[:10], 1):
                dividend_str = f"是({stock['dividend_count']})" if stock['has_dividend'] else "否"
                print(f"{i:<4} {stock['code']:<8} {stock['name']:<12} "
                      f"{stock['current_price']:<8.2f} {stock['market_cap']:<8.1f} "
                      f"{stock['pe_ratio']:<6.1f} {dividend_str:<6} {stock['screening_score']:<6.1f}")
                      
            # 推荐前5名
            print(f"\n🎯 推荐候选股票 (前5名):")
            for i, stock in enumerate(screened_stocks[:5], 1):
                print(f"   {i}. {stock['code']} {stock['name']} - 评分: {stock['screening_score']:.1f}")
                
        else:
            print(f"\n⚠️ 没有股票通过筛选条件，建议调整筛选标准")
            
    def save_screening_results(self, screened_stocks, filename="screening_results.csv"):
        """保存筛选结果"""
        if screened_stocks:
            df = pd.DataFrame(screened_stocks)
            df.to_csv(filename, index=False, encoding='utf-8-sig')
            print(f"\n💾 筛选结果已保存到 {filename}")
            return filename
        return None

if __name__ == "__main__":
    # 创建筛选器实例
    screener = AdvancedStockScreener()
    
    # 执行筛选
    results = screener.screen_candidate_stocks()
    
    # 保存结果
    screener.save_screening_results(results)
    
    print(f"\n🎉 第二步：股票筛选系统执行完成！")
    print(f"   • 成功筛选出 {len(results)} 只符合条件的候选股票")
    print(f"   • 准备进入第三步：技术分析系统")
